

#include "lua_compile.h"

/********************************************************************************/

static int dump_writer(lua_State* L, const void* p, size_t sz, void* ud) {
  std::string* ps = (std::string*)ud;
  return ps->append((char*)p, sz), LUA_OK;
}

static bool encrypt(lua_State* L, std::string& data) {
  int top = lua_gettop(L);
  lua_getglobal(L,     "crypto");
  lua_getfield (L, -1, "aes");
  lua_getfield (L, -1, "encrypt");
  lua_pushlstring(L, data.c_str(), data.size());
  lua_pushliteral(L, LUAOS_AES_KEY);
  if (lua_pcall(L, 2, 1, 0) != LUA_OK) {
    return false;
  }
  size_t size = 0;
  const char* p = luaL_checklstring(L, -1, &size);
  data.assign(p, size);
  lua_settop(L, top);
  char* p1 = (char*)data.c_str();
  const char* p2 = "\x55\xAA";
  for (size_t i = 0; i < size; i++) {
    *p1 ^= p2[i & 1];
  }
  return true;
}

static int luac_compile(lua_State* L) {
  std::string data;
  const char* infile  = luaL_checkstring(L, 1);
  const char* outfile = luaL_optstring(L, 2, nullptr);
  int result = luaL_loadfile(L, infile);
  if (result == LUA_OK) {
    result = lua_dump(L, dump_writer, &data, 0);
  }
  if (result != LUA_OK) {
    outfile ? lua_pushboolean(L, 0) : lua_pushnil(L);
    return 1;
  }
  if (encrypt(L, data) == false) {
    outfile ? lua_pushboolean(L, 0) : lua_pushnil(L);
    return 1;
  }
  if (outfile == nullptr) {
    lua_pushlstring(L, data.c_str(), data.size());
    return 1;
  }
  FILE* fpw = fopen(outfile, "wb");
  if (fpw == nullptr) {
    lua_pushboolean(L, 0);
    return 1;
  }
  fwrite(data.c_str(), 1, data.size(), fpw);
  fclose(fpw);
  lua_pushboolean(L, 1);
  return 1;
}

/********************************************************************************/

LUAOS_API int luaopen_compile(lua_State* L) {
  const luaL_Reg methods[] = {
    { "compile",  luac_compile  }, /* os.compile(inname, outname) */
    { NULL,       NULL          }
  };
  return new_module(L, "os", methods);
}

/********************************************************************************/
